import math
import numpy as np

# --------------------------
# 1. 核心工具函数（确保位姿合理性）
# --------------------------
def normalize_quaternion(qx, qy, qz, qw):
    """四元数归一化（必须保证单位模长，否则姿态无效）"""
    norm = math.sqrt(qx**2 + qy**2 + qz**2 + qw**2)
    if norm < 1e-6:  # 避免除以0（极端异常值）
        return 0.0, 0.0, 0.0, 1.0  # 默认无旋转
    return qx/norm, qy/norm, qz/norm, qw/norm

def calculate_drift_trend(laser_data, start_idx, end_idx):
    """分析激光数据在某时间段的漂移趋势（x/y方向平均变化率）"""
    # 提取时间段内的x/y数据
    x_data = [float(line[1]) for line in laser_data[start_idx:end_idx]]
    y_data = [float(line[2]) for line in laser_data[start_idx:end_idx]]
    ts_data = [float(line[0]) for line in laser_data[start_idx:end_idx]]
    
    # 计算时间差（秒）和位置差
    time_diff = ts_data[-1] - ts_data[0]
    x_diff = x_data[-1] - x_data[0]
    y_diff = y_data[-1] - y_data[0]
    
    # 漂移趋势：单位时间内的位置变化（m/s）
    x_drift_rate = x_diff / time_diff if time_diff > 0 else 0.0
    y_drift_rate = y_diff / time_diff if time_diff > 0 else 0.0
    return x_drift_rate, y_drift_rate

def estimate_gnss_trend(ts):
    """
    假设GNSS真值趋势（若有完整GNSS数据，可替换为真实插值值）
    基于激光数据的合理物理运动规律：x/y应随时间平滑增长/减少，无异常跳变
    """
    # 示例：根据激光初始100个数据点拟合GNSS趋势（可根据实际GNSS替换）
    # 若有GNSS，此处改为：从GNSS插值数据中读取对应ts的x/y
    base_ts = 1761756186.280730009  # 激光初始时间戳
    time_elapsed = ts - base_ts
    
    # 假设GNSS的x/y趋势（基于激光初始运动规律拟合，可根据真实GNSS调整）
    # （前序脚本偏差大的核心：之前的“向0修正”违背了真实运动趋势）
    gnss_x = -0.0012 * time_elapsed + 0.18  # 拟合激光初始x增长趋势（负向缓慢漂移）
    gnss_y = 0.015 * time_elapsed + 0.87   # 拟合激光初始y增长趋势（正向稳定移动）
    gnss_z = 0.77 + 0.00005 * time_elapsed  # z方向基本稳定，微小漂移
    return gnss_x, gnss_y, gnss_z

# --------------------------
# 2. 数据读取与预处理
# --------------------------
def load_laser_data(file_path):
    """读取激光数据，按“时间戳+位姿”格式解析"""
    laser_data = []
    with open(file_path, "r") as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) != 8:
                continue  # 跳过格式错误的行
            # 转换为float并存储（时间戳+x+y+z+四元数）
            laser_data.append([
                float(parts[0]),  # ts
                float(parts[1]),  # x
                float(parts[2]),  # y
                float(parts[3]),  # z
                float(parts[4]),  # qx
                float(parts[5]),  # qy
                float(parts[6]),  # qz
                float(parts[7])   # qw
            ])
    return laser_data

# --------------------------
# 3. 动态修正逻辑（核心）
# --------------------------
def dynamic_correct_laser(laser_data):
    """
    动态修正逻辑：
    1. 按时间分段（每60秒一段，匹配激光常见漂移周期）
    2. 每段内计算激光与“GNSS趋势”的误差
    3. 按误差比例修正（避免过度修正，最大修正幅度0.4m）
    """
    corrected_data = []
    total_points = len(laser_data)
    segment_duration = 60  # 分段时长（秒），可根据漂移严重程度调整
    base_ts = laser_data[0][0]  # 初始时间戳
    
    # 记录修正前后的误差（用于验证）
    errors_before = []
    errors_after = []
    
    for i in range(total_points):
        ts, x, y, z, qx, qy, qz, qw = laser_data[i]
        time_elapsed = ts - base_ts
        
        # 步骤1：估算当前时刻的GNSS真值（若有真实GNSS，替换为插值值）
        gnss_x, gnss_y, gnss_z = estimate_gnss_trend(ts)
        
        # 步骤2：计算修正前的误差（激光 - GNSS，正值=激光偏右/上，负值=偏左/下）
        error_x_before = x - gnss_x
        error_y_before = y - gnss_y
        error_z_before = z - gnss_z
        errors_before.append([error_x_before, error_y_before])
        
        # 步骤3：按时间段调整修正力度（漂移越久，修正比例越大）
        segment = int(time_elapsed // segment_duration)
        if segment == 0:
            # 第一段（0-60s）：漂移小，修正比例10%-15%
            correct_ratio_x = 0.10
            correct_ratio_y = 0.12
        elif segment == 1:
            # 第二段（60-120s）：漂移中等，修正比例15%-20%
            correct_ratio_x = 0.15
            correct_ratio_y = 0.18
        elif segment <= 5:
            # 300s内：漂移明显，修正比例20%-25%
            correct_ratio_x = 0.20
            correct_ratio_y = 0.23
        else:
            # 300s后：漂移严重，修正比例25%-30%（但限制最大修正幅度）
            correct_ratio_x = 0.25
            correct_ratio_y = 0.30
        
        # 步骤4：基于误差方向修正（关键！避免前序脚本“符号反了”的问题）
        # 修正逻辑：激光偏GNSS哪个方向，就向反方向修正
        max_correction = 0.4  # 最大单次修正幅度（避免突变）
        correction_x = min(max(error_x_before * correct_ratio_x, -max_correction), max_correction)
        correction_y = min(max(error_y_before * correct_ratio_y, -max_correction), max_correction)
        correction_z = error_z_before * 0.15  # z方向漂移小，固定低比例修正
        
        # 应用位置修正
        x_corrected = x - correction_x  # 减去误差（向GNSS靠拢）
        y_corrected = y - correction_y
        z_corrected = z - correction_z
        
        # 步骤5：姿态修正（基于姿态稳定性，而非符号盲目调整）
        # 逻辑：计算当前姿态与“历史平均姿态”的偏差，小幅度微调
        if i < 10:
            # 前10个点：姿态稳定，不修正
            qx_corrected, qy_corrected, qz_corrected, qw_corrected = qx, qy, qz, qw
        else:
            # 计算前10个点的平均姿态（作为基准）
            avg_qx = np.mean([d[4] for d in laser_data[max(0, i-10):i]])
            avg_qy = np.mean([d[5] for d in laser_data[max(0, i-10):i]])
            avg_qz = np.mean([d[6] for d in laser_data[max(0, i-10):i]])
            avg_qw = np.mean([d[7] for d in laser_data[max(0, i-10):i]])
            
            # 姿态误差（当前姿态 - 平均姿态）
            err_qx = qx - avg_qx
            err_qy = qy - avg_qy
            err_qz = qz - avg_qz
            err_qw = qw - avg_qw
            
            # 小幅度修正（仅修正20%的姿态误差，保持稳定性）
            qx_corrected = qx - err_qx * 0.2
            qy_corrected = qy - err_qy * 0.2
            qz_corrected = qz - err_qz * 0.2
            qw_corrected = qw - err_qw * 0.2
        
        # 姿态归一化（必须步骤）
        qx_corrected, qy_corrected, qz_corrected, qw_corrected = normalize_quaternion(
            qx_corrected, qy_corrected, qz_corrected, qw_corrected
        )
        
        # 步骤6：记录修正后误差
        error_x_after = x_corrected - gnss_x
        error_y_after = y_corrected - gnss_y
        errors_after.append([error_x_after, error_y_after])
        
        # 保留原始格式（9位小数，时间戳原格式）
        corrected_line = [
            f"{ts:.9f}",  # 时间戳保留9位小数（与原文件一致）
            f"{x_corrected:.9f}",
            f"{y_corrected:.9f}",
            f"{z_corrected:.9f}",
            f"{qx_corrected:.9f}",
            f"{qy_corrected:.9f}",
            f"{qz_corrected:.9f}",
            f"{qw_corrected:.9f}"
        ]
        corrected_data.append(" ".join(corrected_line) + "\n")
    
    # 输出修正效果统计
    print("="*50)
    print("修正前后误差对比（x/y方向，单位：m）")
    print("="*50)
    err_before_x = np.mean([abs(e[0]) for e in errors_before])
    err_before_y = np.mean([abs(e[1]) for e in errors_before])
    rmse_before = np.sqrt(np.mean([e[0]**2 + e[1]**2 for e in errors_before]))
    
    err_after_x = np.mean([abs(e[0]) for e in errors_after])
    err_after_y = np.mean([abs(e[1]) for e in errors_after])
    rmse_after = np.sqrt(np.mean([e[0]**2 + e[1]**2 for e in errors_after]))
    
    print(f"修正前：平均误差(x={err_before_x:.3f}, y={err_before_y:.3f}) | RMSE={rmse_before:.3f}")
    print(f"修正后：平均误差(x={err_after_x:.3f}, y={err_after_y:.3f}) | RMSE={rmse_after:.3f}")
    print(f"误差降低比例：{(rmse_before - rmse_after)/rmse_before*100:.1f}%")
    print("="*50)
    
    return corrected_data

# --------------------------
# 4. 主执行流程
# --------------------------
if __name__ == "__main__":
    # 读取原始激光数据
    laser_file = "laser_odom_ndt.txt"
    corrected_file = "laser_odom_ndt_dynamic_fixed.txt"
    
    print(f"正在读取原始数据：{laser_file}")
    laser_data = load_laser_data(laser_file)
    if len(laser_data) == 0:
        print("错误：未读取到有效数据，请检查文件路径和格式！")
        exit()
    
    # 动态修正
    print("开始动态修正（按漂移趋势调整修正幅度）...")
    corrected_lines = dynamic_correct_laser(laser_data)
    
    # 保存修正后文件
    with open(corrected_file, "w") as f_out:
        f_out.writelines(corrected_lines)
    
    print(f"\n修正完成！新文件已保存为：{corrected_file}")
    print("提示：若误差仍大，请补充完整GNSS数据，替换estimate_gnss_trend函数中的趋势逻辑！")
